"""
该脚本示例了一个 WebSocket 客户端 + SIM 模拟做市商 的使用
仅根据目标深度和当前深度, 持续向服务器发送报价单或撤单
"""

from __future__ import annotations

import asyncio
import dataclasses
import logging
import random
import sys
import traceback
from typing import Generator, Any, List, Dict, Tuple

import dotenv
import httpx

from datetime import datetime, timedelta
from eztrade_client import WebSocketClient
from eztrade_client import (
    Direction,
    Order,
    OrderStatus,
    PyMessage,
    PositionEffect,
    Portfolio,
    TickData,
    get_logger,
    OrderOperation,
    Operation,
    PriceLevel,
)

dotenv.load_dotenv()  # 从 .env 文件加载环境变量

# %%
SERVER = "127.0.0.1"
PORT = 7878
BASE_URL = f"https://{SERVER}:{PORT}"
STAGE_ID = 2
PATH = f"/ws/{STAGE_ID}"
USER_ID = 0  # 0 代表模拟做市商
STAGE_PAYLOAD = {
    "id": STAGE_ID,
    "stage_type": "SIM",
    "name": f"test-sim-{STAGE_ID}",
    "description": "测试 SIM Stage",
    "sim_cfg": {
        "time_unit_us": 600_000_000,
        "time_multiplier": 600.0,
    },
}
SYMBOL = "au2512"  # 模拟做市商交易的合约代码
logger = get_logger("Simulate", level=logging.INFO)
logger.addFileHandler("logs/sim_maker.log", level=logging.DEBUG)


# %%
@dataclasses.dataclass
class AccountsInfo:
    act_orders_long: Dict[int, Order] = dataclasses.field(default_factory=dict)
    act_orders_short: Dict[int, Order] = dataclasses.field(default_factory=dict)
    done_orders: Dict[int, Order] = dataclasses.field(default_factory=dict)
    portfolio: Portfolio | None = dataclasses.field(default=None)
    lastprice: Dict[str, float] = dataclasses.field(default_factory=dict)

    def get_orders_by_cum(
        self, direction: Direction, price: float, amt: int
    ) -> List[Order]:
        orders = []
        amt_sum = 0
        def sort_orders_by_amt(target_orders: Dict[int, Order]) -> List[Order]:
            return sorted(
                filter(lambda order: order.price == price, target_orders.values()),
                key=lambda order: order.quantity,
            )

        targets = sort_orders_by_amt(
            self.act_orders_long
            if direction == Direction.Long
            else self.act_orders_short
        )
        for order in targets:
            if amt_sum >= amt:
                break
            orders.append(order)
            amt_sum += order.quantity
        return orders


@dataclasses.dataclass
class OrderBookDepth:
    """
    维护一个目标合约的买卖盘深度, 据此进行做市
    """

    symbol: str
    bids: List[PriceLevel] = dataclasses.field(default_factory=list)
    asks: List[PriceLevel] = dataclasses.field(default_factory=list)

    def __post_init__(self):
        self.bids.sort(key=lambda x: x.price, reverse=True)
        self.asks.sort(key=lambda x: x.price)

    @staticmethod
    def process_side(depth_target: List[PriceLevel], depth_current: List[PriceLevel]):
        d_target = {lvl.price: lvl.volume for lvl in depth_target}
        d_current = {lvl.price: lvl.volume for lvl in depth_current}
        price_union = {lvl.price for lvl in depth_target}.union(
            {lvl.price for lvl in depth_current}
        )
        return {
            price: d_target.get(price, 0) - d_current.get(price, 0)
            for price in price_union
        }

    def __sub__(self, other: OrderBookDepth) -> Tuple[List[Order], List[Order]]:
        """
        计算当前盘口到目标盘口的差异, 并生成相应的订单列表
        """
        orders_open = []
        orders_cancel = []
        bids_diff = self.process_side(self.bids, other.bids)
        asks_diff = self.process_side(self.asks, other.asks)
        logger.info(f"Bids diff: {bids_diff}")
        logger.info(f"Asks diff: {asks_diff}")
        for price, diff in bids_diff.items():
            if diff > 0:
                orders_open.append(
                    Order(
                        self.symbol,
                        Direction.Long,
                        PositionEffect.Open,
                        abs(diff),
                        price,
                    )
                )
            elif diff < 0:
                orders_cancel.extend(
                    ACCOUNTS_INFO.get_orders_by_cum(Direction.Long, price, abs(diff))
                )

        for price, diff in asks_diff.items():
            if diff > 0:
                orders_open.append(
                    Order(
                        self.symbol,
                        Direction.Short,
                        PositionEffect.Open,
                        abs(diff),
                        price,
                    )
                )
            elif diff < 0:
                orders_cancel.extend(
                    ACCOUNTS_INFO.get_orders_by_cum(Direction.Short, price, abs(diff))
                )
        return orders_open, orders_cancel


ACCOUNTS_INFO = AccountsInfo()
TARGET_DEPTH = OrderBookDepth(
    SYMBOL,
    [PriceLevel(price=4995.0, volume=10), PriceLevel(price=4990.0, volume=20)],
    [PriceLevel(price=5005.0, volume=10), PriceLevel(price=5010.0, volume=20)],
)  # 实际目标盘口应该是动态的


# %%
def fmt_from_timestamp_us(timestamp_us: int) -> str:
    dt = datetime.fromtimestamp(timestamp_us / 1_000_000)
    return dt.strftime("%Y-%m-%d %H:%M:%S.%f")


def maker_generator(client: WebSocketClient) -> Generator[Any, Any]:
    """
    模拟做市商：获取最新TickData, 并持续向服务器发送报价单。
    更精确的方案是依据最新的TickData, 匹配目标OrderBookDepth, 进行最小变动
    """
    while True:
        if ACCOUNTS_INFO.portfolio is None:
            logger.warning("Portfolio 信息尚未获取, 无法进行做市操作")
            _ = yield False
            continue
        _tick_data_pre: TickData = yield True
        depth_now = OrderBookDepth(SYMBOL, _tick_data_pre.bids, _tick_data_pre.asks)
        time_fmt = fmt_from_timestamp_us(_tick_data_pre.timestamp_us)
        logger.info(f"[{time_fmt}] Depth: {depth_now}")
        orders_open, orders_cancel = TARGET_DEPTH - depth_now
        for order in orders_open:
            msg = PyMessage(order)
            client.send_message(msg)
        for order in orders_cancel:
            msg = PyMessage(OrderOperation(order.id, Operation.Cancel))
            client.send_message(msg)
            logger.info(f"Sent cancel for Order ID: {order.id}")


async def process_events(client: WebSocketClient) -> None:
    """持续消耗数据队列, 并根据不同类型的事件执行相应操作。"""
    order_generator = maker_generator(client)
    order_generator.send(None)  # 初始化生成器
    while True:
        if client.buffer:
            data = client.buffer.pop() # 取出最新数据
            if len(client.buffer) > 10:
                logger.warning("Size of buffer: %s", len(client.buffer))
            try:
                message = PyMessage.from_binary_bytes(data)
            except Exception as exc:  # noqa: BLE001
                logger.error("Failed to decode message: %s, raw: %s", exc, data)
                continue

            match message.get_data_type():
                case "Authority":
                    auth = message.get_data()
                    logger.debug("UserId: %s", auth.client_id)
                    logger.info("Query current positions")
                    client.send_message(PyMessage.query_account())
                    client.send_message(PyMessage.subscribe_tick(SYMBOL))
                case "TickData":
                    tickdata: TickData = message.get_data()
                    ACCOUNTS_INFO.lastprice[tickdata.symbol] = tickdata.last
                    # time_fmt = fmt_from_timestamp_us(tickdata.timestamp_us)
                    # logger.info(f"[{time_fmt}] Received TickData: {tickdata}")
                    order_generator.send(tickdata)
                case "Order":
                    order: Order = message.get_data()
                    time_fmt = fmt_from_timestamp_us(order.timestamp_us)
                    if order.status == OrderStatus.Unfilled:
                        logger.debug(f"[{time_fmt}] Order {order} has been sent.")
                        if order.direction == Direction.Long:
                            ACCOUNTS_INFO.act_orders_long.update({order.id: order})
                        else:
                            ACCOUNTS_INFO.act_orders_short.update({order.id: order})
                    elif order.status == OrderStatus.PartiallyFilled:
                        logger.info(f"[{time_fmt}] Order {order} has been partially filled.")
                        if order.direction == Direction.Long:
                            ACCOUNTS_INFO.act_orders_long.update({order.id: order})
                        else:
                            ACCOUNTS_INFO.act_orders_short.update({order.id: order})
                    elif order.status in [
                        OrderStatus.Filled,
                        OrderStatus.PartiallyFilled,
                        OrderStatus.Canceled,
                        OrderStatus.Rejected,
                    ]:
                        logger.debug(f"[{time_fmt}] Order {order} has finished!")
                        ACCOUNTS_INFO.act_orders_long.pop(order.id, None)
                        ACCOUNTS_INFO.act_orders_short.pop(order.id, None)
                        ACCOUNTS_INFO.done_orders.update({order.id: order})
                    else:
                        pass
                case "Portfolio":
                    portfolio: Portfolio = message.get_data()
                    ACCOUNTS_INFO.portfolio = portfolio
                case "String":
                    text: str = message.get_data()
                    logger.info("Received string data: %s", text)
                case _:
                    logger.info("Received data: %s", message)
        else:
            await asyncio.sleep(0.001)

        if client.stop_flag.is_set():
            logger.info("客户端接收停止标志已设置, 退出消费者循环")
            break


async def main() -> None:
    async with httpx.AsyncClient(
        base_url=BASE_URL,
        timeout=10.0,
        verify=False,
        trust_env=False,
    ) as async_client:
        resp = await async_client.post("/stages/create", json=STAGE_PAYLOAD)
        assert resp.status_code in [409, 201], f"Create stage failed: {resp.text}"

    client = WebSocketClient()
    client.create(SERVER, PORT, PATH, USER_ID)
    try:
        async with asyncio.TaskGroup() as tg:
            tg.create_task(process_events(client))
    except ExceptionGroup as exc_group:
        for exc in exc_group.exceptions:
            logger.error(
                "主循环异常退出: %s: %s", exc, "".join(traceback.format_exception(exc))
            )
    except asyncio.CancelledError:
        logger.info("主循环取消退出")
    finally:
        client.stop_loop_thread()

    async with httpx.AsyncClient(
        base_url=BASE_URL,
        timeout=10.0,
        verify=False,
        trust_env=False,
    ) as async_client:
        resp = await async_client.delete(f"/stages/{STAGE_ID}")
        if resp.status_code != 204:
            logger.error("Delete stage failed: %s", resp.text)
        else:
            logger.info("Stage %s deleted successfully.", STAGE_ID)

    sys.exit(1)


# %%
if __name__ == "__main__":
    asyncio.run(main())
